import numpy as np

import preprocessingfile
import train
import model
def align_columns(x_train, x_test):
    # 获取列名集合
    train_columns = set(x_train.columns)
    test_columns = set(x_test.columns)

    # 找出 x_train 中缺少的列
    missing_in_train = test_columns - train_columns
    for col in missing_in_train:
        x_train[col] = 0  # 添加缺失列，值设为 0

    # 找出 x_test 中缺少的列
    missing_in_test = train_columns - test_columns
    for col in missing_in_test:
        x_test[col] = 0  # 添加缺失列，值设为 0

    # 按列名排序，确保列的顺序一致
    x_train = x_train[sorted(x_train.columns)]
    x_test = x_test[sorted(x_test.columns)]

    return x_train, x_test

if __name__ == '__main__':
    # 测试数据预处理
    x_train, y_train = preprocessingfile.preprocessor_train_data(['data/train/train.csv','data/train/X_train.csv'])
    x_test = preprocessingfile.preprocessor_test_data(['data/test/test.csv','data/test/X_test.csv'])
    x_train, x_test = align_columns(x_train, x_test)
    w, b = train.logistic_regression(x_train, y_train, 32, 100, 0.01, True)
    weights = {'w': w, 'b': b}
    model.save_model(weights)
    output_file = 'output.csv'
    x_test = np.nan_to_num(x_test.values, nan=0.0)
    result = train.infer(x_test, w, b)
    with open(output_file, 'w') as f:
        f.write("id,label\n")
        for index,value in enumerate(result):
            f.write("%d,%d\n" % (index+1,value.item()))